#include "pointcloud_tools.h"
#include "point_preprocess.h"
#include <iostream>
#include "point_common.h"
#include <fstream>
using namespace pointcloud_tools;
// using namespace PointPreprocess;
// using export_blocks_to_files = PointPreprocess::export_blocks_to_files;

// 八叉树递归分割函数
void recursive_split(int z_low, int y_low, int x_low, int z_high, int y_high, int x_high,
                     std::vector<std::vector<int>>& pt_mapping,
                     std::vector<BlockConfig>& bcfg_output,
                     BlockConfig& bcfg_global,
                     std::vector<std::vector<int>>& valid_filter,
                     uint64_t& current_pt_addr, uint64_t& current_in_addr, uint64_t& current_rs_addr,
                     int ich_loop_num, int och_loop_num,
                     uint8_t* search_range, int block_size = 1024) {
    // Padding
    int pad_z = search_range[0];
    int pad_y = search_range[1];
    int pad_x = search_range[2];

    int z_low_pad = z_low - pad_z;
    int y_low_pad = y_low - pad_y;
    int x_low_pad = x_low - pad_x;
    int z_high_pad = z_high + pad_z;
    int y_high_pad = y_high + pad_y;
    int x_high_pad = x_high + pad_x;

    // std::cout<<"Number of points befor split: "<<bcfg_global.pt_data.size()<<std::endl;

    std::vector<int> indices_pad;
    // Find the indices of points within the padded range
    for (size_t i = 0; i < bcfg_global.pt_data.size(); i++) {
        const auto& coord = bcfg_global.pt_data[i];
        // std::cout<<"coord: "<<coord<<std::endl;
        // Note: COMPARING should be done with both signed types!
        int64_t x = coord & 0xFFF;
        int64_t y = (coord >> 12) & 0xFFF;
        int64_t z = (coord >> 24) & 0xFFF;
        // std::cout<<"x,y,z: "<<x<<","<<y<<","<<z<<std::endl;
        // std::cout<<"padded range:(x, y, z) "<<x_low_pad<<","<<x_high_pad<<","<<y_low_pad<<","<<y_high_pad<<","<<z_low_pad<<","<<z_high_pad<<std::endl;
        if (x >= x_low_pad && x <= x_high_pad &&
            y >= y_low_pad && y <= y_high_pad &&
            z >= z_low_pad && z <= z_high_pad) {
            indices_pad.push_back(i);
            // std::cout<<"True"<<std::endl;
        }
    }

    // std::cout<<"Number of points wirhin the padded range: "<<indices_pad.size()<<std::endl;

    if (indices_pad.empty()) {
        return; // No points in this range, return
    }

    if (indices_pad.size() <= block_size) {
        // Create a block config 
        // Create a filter for this block
        BlockConfig current_bcfg;
        std::vector<int> current_filter;

        pt_mapping.push_back(indices_pad);

        current_bcfg.BASE_PT_ADDR = current_pt_addr; // 当前块的坐标起始地址
        current_bcfg.BASE_IN_ADDR = current_in_addr;
        current_bcfg.BASE_RS_ADDR = current_rs_addr;
        current_bcfg.cfg_cam_valid_num = indices_pad.size();
        current_bcfg.cfg_total_point_num = indices_pad.size();
        // std::cout<<"Block size: "<<indices_pad.size()<<std::endl;

        int count = 0;
        for (int loop = 0; loop < ich_loop_num + 1; loop++) {
            for (int idx : indices_pad) {
                if (loop == 0) {
                    current_bcfg.pt_data.push_back(bcfg_global.pt_data[idx]);
                    int64_t x = bcfg_global.pt_data[idx] & 0xFFF;
                    int64_t y = (bcfg_global.pt_data[idx] >> 12) & 0xFFF;
                    int64_t z = (bcfg_global.pt_data[idx] >> 24) & 0xFFF;
                    if (x >= x_low && x <= x_high &&
                        y >= y_low && y <= y_high &&
                        z >= z_low && z <= z_high) {
                        // current_filter.push_back(idx);
                        current_filter.push_back(count);
                        // if (idx == 0) {
                        //     std::cout<<"The #0 point belongs to block:"<<bcfg_output.size()<<std::endl;
                        // }
                    }
                    count++;
                }
                current_bcfg.in_data.push_back(bcfg_global.in_data[idx + loop * bcfg_global.pt_data.size()]);
            }
        }
        
        // 计算下一个块的地址，要求512bit（64字节，0x40）对齐
        // 计算下一个块的地址，要求4K 4*1024*8（64字节，0x1000）对齐
        #ifdef SIMULATION
            auto align64B = [](uint64_t addr) {
                return (addr + 0xFFF) & ~(uint64_t)(0xFFF);
            };
            // auto align64B = [](uint64_t addr) {
            //     return (addr + 0xFF) & ~(uint64_t)(0xFF);
            // };
            current_pt_addr = align64B(current_pt_addr + indices_pad.size() * sizeof(int64_t));
            current_in_addr = align64B(current_in_addr + indices_pad.size() * sizeof(int64_t) * (1 + ich_loop_num));
            current_rs_addr = align64B(current_rs_addr + indices_pad.size() * sizeof(int64_t) * (1 + och_loop_num));
        #endif

        bcfg_output.push_back(current_bcfg);
        valid_filter.push_back(current_filter); // current_filter中点的索引 = 输入中点的索引 = 结果中点的索引 (SubMCONV)
    }

    else {
        // std::cout<<"Block size out of range:: "<<indices_pad.size()<<std::endl;
        // Split the range into 8 sub-blocks
        int z_mid = (z_low + z_high) / 2;
        int y_mid = (y_low + y_high) / 2;
        int x_mid = (x_low + x_high) / 2;

        for (int i = 0; i < 8; i++) {
            int new_z_low = (i / 4) % 2 == 0 ? z_low : z_mid + 1;
            int new_z_high = (i / 4) % 2 == 0 ? z_mid : z_high;
            int new_y_low = (i / 2) % 2 == 0 ? y_low : y_mid + 1;
            int new_y_high = (i / 2) % 2 == 0 ? y_mid : y_high;
            int new_x_low = i % 2 == 0 ? x_low : x_mid + 1;
            int new_x_high = i % 2 == 0 ? x_mid : x_high;

            recursive_split(new_z_low, new_y_low, new_x_low, new_z_high, new_y_high, new_x_high,
                            pt_mapping, bcfg_output, bcfg_global, valid_filter,
                            current_pt_addr, current_in_addr, current_rs_addr,
                            ich_loop_num, och_loop_num,
                            search_range, block_size);
        }
    }
}

// 八叉树分割代码，直接传入整个点云完整块进行分割
void octree_partition(BlockConfig& bcfg, const int block_size, std::vector<std::vector<int>>& pt_mapping,
                        std::vector<BlockConfig>& bcfg_output, std::vector<std::vector<int>>& valid_filter,
                        uint8_t* search_range, uint16_t* spatial_shape,
                        int ich_loop_num, int och_loop_num) {
    // bcfg: input after voxelization
    // bcfg_output: output block config
    // block_size: max number of points in a block, default 1024
    // spatial_shape: [z, y, x] defalut 0~z, 0~y, 0~x

    // Spatial shape before split
    uint16_t vox_x_min = spatial_shape[2];
    uint16_t vox_y_min = spatial_shape[1];
    uint16_t vox_z_min = spatial_shape[0];

    uint16_t vox_x_max = spatial_shape[5];
    uint16_t vox_y_max = spatial_shape[4];
    uint16_t vox_z_max = spatial_shape[3];

    recursive_split(vox_z_min, vox_y_min, vox_x_min, vox_z_max, vox_y_max, vox_x_max, 
                    pt_mapping, bcfg_output, bcfg, valid_filter,
                    bcfg.BASE_PT_ADDR, bcfg.BASE_IN_ADDR, bcfg.BASE_RS_ADDR, 
                    ich_loop_num, och_loop_num,
                    search_range, block_size);
}

void merge_RESULT_feat(const std::vector<BlockConfig>& bcfg_partition, int valid_pt_num, int och_loop_num,
                        const std::vector<std::vector<int>>& pt_mapping, const std::vector<std::vector<int>>& valid_filter, std::vector<uint64_t>& merged_res) {
    // Fisrt initialize the merged_res vector
    merged_res.resize(valid_pt_num * (1 + och_loop_num), 0);
    
    for (int i = 0; i < bcfg_partition.size(); i++) {
        // For each block, get the rs_data and pt_mapping
        const BlockConfig& bcfg_tmp = bcfg_partition[i];
        const std::vector<uint64_t>& rs_tmp = bcfg_tmp.rs_data;
        const std::vector<int>& filter_tmp = valid_filter[i];
        const std::vector<int>& pt_mapping_tmp = pt_mapping[i];
        // assert(rs_tmp.size() == bcfg_tmp.cfg_cam_valid_num * (1 + och_loop_num) && ""
        //        "The size of rs_tmp is not equal to cfg_cam_valid_num * (1 + och_loop_num)");
        // if (rs_tmp.size() != bcfg_tmp.cfg_cam_valid_num * (1 + och_loop_num)) {
        //     std::cerr << "Error: The size of rs_tmp is not equal to cfg_cam_valid_num * (1 + och_loop_num)" << std::endl;
        //     continue; // Skip this block if the sizes do not match
        // }
        for(int k = 0; k < och_loop_num + 1; k++) {
            for (int j = 0; j < filter_tmp.size(); j++) {
                int pt_idx = filter_tmp[j];
                assert(pt_idx < bcfg_tmp.cfg_cam_valid_num);
                int pt_mapping_idx = pt_mapping_tmp[pt_idx];            
                merged_res[pt_mapping_idx + k * valid_pt_num] = rs_tmp[pt_idx + k * bcfg_tmp.cfg_cam_valid_num];
            }
        }   
    }

    std::cout << "Merged result done!" << std::endl;
    // return merged_res;
}

// 将合并结果写入文本文件，每行代表一个点的所有通道数据
void export_merged_results(const std::vector<uint64_t>& merged_res, int valid_pt_num, int och_loop_num, const std::string& filename) {
    // 打开输出文件
    std::ofstream outfile(filename);
    if (!outfile.is_open()) {
        std::cerr << "无法打开文件: " << filename << std::endl;
        return;
    }

    // 创建union用于将uint64_t拆分为8个int8_t
    union data64to8 {
        uint64_t in64;
        int8_t out8[8];
    } converter;

    // 遍历每个点
    for (int pt_idx = 0; pt_idx < valid_pt_num; pt_idx++) {
        std::string line = "";
        // 遍历每个loop
        for (int k = 0; k <= och_loop_num; k++) {
            // 获取当前点在当前loop的数据
            uint64_t data = merged_res[pt_idx + k * valid_pt_num];
            converter.in64 = data;
            
            // 将数据拆分为8个int8_t并添加到行中
            for (int j = 0; j < 8; j++) {
                if (!line.empty()) {
                    line += ",";
                }
                line += std::to_string(static_cast<int>(converter.out8[j]));
            }
        }
        // 写入一行
        outfile << line << std::endl;
    }
    
    outfile.close();
    std::cout << "MERGED RESULT EXPORT DONE: " << filename << std::endl;
}

int main(int argc, char** argv) {
    Timer timer, timer_preprocess, timer_wt_init, timer_run;
    timer.reset();
    timer_preprocess.reset();
    timer_wt_init.reset();
    timer_run.reset();
    timer.start();

    #ifdef SIMULATION
        Pointcloud_utils ptu;   
        ptu.ddrSize = 8 * 8 * 1024 * 1024 * sizeof(int); //256MB
        assert(ptu.Model_DDR_alloc("DDR_32MB.bin"));
        printf("Simulation Mode is ON!\n");
    #else
        Pointcloud_online_utils ptu;
        ptu.pmem_mmap_init();
        printf("Onchip Mode is ON!\n");
        // ptu.sram_rw_test();
        // ptu.ddr_rw_test();
    #endif
    ptu.DDR_BASE_ADDR = 0x95000000;

    /////////////////////////////////////
    // layer1                          //
    // SUBCONV: ICH = 4, OCH = 32      //
    /////////////////////////////////////
    timer_preprocess.start();
    ptu.layer_config.cfg_inch_loop_num = 0;
    ptu.layer_config.cfg_core_enable = 0xf; // 32 output channels
    ptu.layer_config.cfg_shift_num = 8;
    ptu.layer_config.BASE_WT_ADDR = ptu.DDR_BASE_ADDR;
    size_t dlen = readBinToVec(ptu.layer_config.wt_data, "../data/golden_data/conv1_weight_packed_pad.bin");
    assert(dlen > 0);
    BlockConfig bcfg;
    dlen = readBinToVec(bcfg.pt_data, "../data/golden_data/voxel_coords.bin");
    assert(dlen > 0);
    dlen = readBinToVec(bcfg.in_data, "../data/golden_data/voxel_features.bin");
    assert(dlen > 0);



    uint16_t valid_pt_num = bcfg.pt_data.size();
    bcfg.cfg_cam_valid_num = valid_pt_num;
    bcfg.cfg_total_point_num = valid_pt_num;
    bcfg.BASE_PT_ADDR = ptu.DDR_BASE_ADDR + 0x01000000;
    bcfg.BASE_IN_ADDR = ptu.DDR_BASE_ADDR + 0x02000000;
    bcfg.BASE_RS_ADDR = ptu.DDR_BASE_ADDR + 0x03000000;

    std::vector<std::vector<int>> pt_mapping;
    std::vector<BlockConfig> bcfg_partition;
    std::vector<std::vector<int>> valid_filter;
    uint16_t spatial_shape[6] = {0, 0, 0, 800, 600, 800}; // [z_min, y_min, x_min, z_max, y_max, x_max]

    // Partition the point cloud into blocks
    octree_partition(bcfg, 1024, 
                            pt_mapping, bcfg_partition, valid_filter,
                            ptu.layer_config.cfg_cam_search_range, 
                            spatial_shape, ptu.layer_config.cfg_inch_loop_num, 3);

    PointPreprocess::export_blocks_to_files(
        bcfg_partition, pt_mapping, valid_filter,
        "output_blocks",
        bcfg.pt_data, bcfg.in_data);

    ptu.block_config.clear();
    for (int i = 0; i < bcfg_partition.size(); i++) { //bcfg_partition.size()
        BlockConfig bcfg_tmp = bcfg_partition[i];
        ptu.block_config.push_back(bcfg_tmp);
    }

    timer_preprocess.stop();
    timer_preprocess.print("Preprocess time");

    timer_wt_init.start();
    ptu.wt_init();
    timer_wt_init.stop();
    timer_wt_init.print("Weight init time for layer1");
    #ifdef SIMULATION
        ptu.compute_init();
        ptu.run();
    #else
        timer_run.start();
        ptu.run_block();
        timer_run.stop();
        timer_run.print("Run time for layer1");
    #endif


    // Filter and merge the results
    std::vector<uint64_t> layer1_merged_res;
    merge_RESULT_feat(ptu.block_config, bcfg.cfg_cam_valid_num, 3, pt_mapping, valid_filter, layer1_merged_res);
    #ifdef SIMULATION
        std::string layer1_output_file = "layer1_output.txt";
        export_merged_results(layer1_merged_res, valid_pt_num, 3, layer1_output_file);    
    #endif
    //////////////////////////////////////
    // layer2                           //
    // SubConv: ich = 32, och = 8       //
    //////////////////////////////////////

    ptu.layer_config.cfg_inch_loop_num = 3;
    ptu.layer_config.cfg_core_enable = 0x1; // 32 output channels
    ptu.layer_config.cfg_shift_num = 8;
    ptu.layer_config.BASE_WT_ADDR = ptu.DDR_BASE_ADDR;
    ptu.layer_config.layer_id = 2; // Layer ID for layer2
    dlen = readBinToVec(ptu.layer_config.wt_data, "../data/golden_data/conv2_weight_packed_pad.bin");
    assert(dlen > 0);
    
    uint64_t DDR_IN_ADDR = ptu.block_config[0].BASE_IN_ADDR;
    uint64_t DDR_RS_ADDR = ptu.block_config[ptu.block_config.size() - 1].BASE_RS_ADDR;
    // clear the in_data for every block
    for (int block_idx = 0; block_idx < ptu.block_config.size(); block_idx++) {
        ptu.block_config[block_idx].in_data.clear();
    }
    // Update the in_data for every block
    for (int block_idx = 0; block_idx < ptu.block_config.size(); block_idx++) {
        int pt_num = ptu.block_config[block_idx].cfg_cam_valid_num;
        std::vector<int> pt_mapping_tmp = pt_mapping[block_idx];
        ptu.block_config[block_idx].in_data.reserve(pt_num * (ptu.layer_config.cfg_inch_loop_num + 1));
        for (int j = 0; j < ptu.layer_config.cfg_inch_loop_num + 1; j++) {
            for (int k = 0; k < pt_num; k++) {
                // if (k >= pt_mapping_tmp.size()) {
                //     std::cerr << "Error: k >= pt_mapping_tmp.size()" << std::endl;
                //     continue;
                // }
                int idx_tmp = pt_mapping_tmp[k] + j * bcfg.cfg_cam_valid_num;
                ptu.block_config[block_idx].in_data.push_back(layer1_merged_res[idx_tmp]);
            }
        }
        ptu.block_config[block_idx].BASE_IN_ADDR = DDR_IN_ADDR;
        ptu.block_config[block_idx].BASE_RS_ADDR = DDR_RS_ADDR;   
        #ifdef SIMULATION
            DDR_IN_ADDR += pt_num * sizeof(uint64_t) * (1 + ptu.layer_config.cfg_inch_loop_num);
            DDR_RS_ADDR += pt_num * sizeof(uint64_t) * (1 + 0);
        #endif
    }

    // int count = 0;
    // for (int i = 0; i < ptu.block_config[0].in_data.size(); i++) {        
    //     printf("in_data[%d],", static_cast<int>(ptu.block_config[0].in_data[i]));       
    //     if (count == 7){
    //         printf("\n");
    //         count = 0;
    //     }
    //     count++;
    // }
    timer_wt_init.reset();
    timer_wt_init.start();
    ptu.wt_init();
    timer_wt_init.stop();
    timer_wt_init.print("Weight init time for layer2");
    #ifdef SIMULATION
        ptu.compute_init();
        ptu.run();
    #else
        timer_run.reset();
        timer_run.start();
        ptu.run_block();
        timer_run.stop();
        timer_run.print("Run time for layer2");
    #endif


    // Filter and merge the results
    std::vector<uint64_t> layer2_merged_res;
    merge_RESULT_feat(ptu.block_config, bcfg.cfg_cam_valid_num, 0, pt_mapping, valid_filter, layer2_merged_res);
    #ifdef DEBUG
        std::string layer2_output_file = "layer2_output.txt";
        export_merged_results(layer2_merged_res, valid_pt_num, 0, layer2_output_file);
    #endif
    timer.stop();
    timer.print("Total time for layer2");

    ///////////////////////////////////////////
    // FC layer                              //
    // ich = 8, och = 2                      //
    ///////////////////////////////////////////    
    float fc_weight[16] = {-0.04190924, 0.08408088, 0.19655997, -0.33468807, 0.35655615, -0.26418057, -0.22881983, 0.03832784,
                0.11310787, 0.27403793, 0.34556514, 0.18055771, 0.13469994, -0.10524856, -0.135478, 0.1271919};

    std::vector<int> voxel_vector;
    for (int idx = 0; idx < bcfg.cfg_cam_valid_num; idx++) {
        float res[2] = {0, 0};
        // uint64_t in_data_tmp = layer2_merged_res[idx];
        Data64to8 in64to8;
        in64to8.in64 = layer2_merged_res[idx];
        for (int j = 0; j < 2; j++) {
            for(int k = 0; k < 8; k++) {
                int8_t ch = in64to8.out8[k];
                res[j] += fc_weight[j * 8 + k] * ch;
            }
        }
        int classID = res[0] > res[1] ? 0 : 1;
        voxel_vector.push_back(classID);
    }

    #ifdef SIMULATION
        ptu.Model_DDR_close(); // Close the simulated DDR file
    #endif

}

  



